import os
import sqlite3
import shutil
from pathlib import Path
import zipfile
import sys
import argparse

def check_and_connect_db(db_path):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row  # 设置返回字典格式的结果
    cursor = conn.cursor()
    return conn, cursor

def zip_dir(folder_to_zip, zip_path):
    """压缩指定文件夹为zip文件"""
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(folder_to_zip):
            for file in files:
                file_path = os.path.join(root, file)
                arcname = os.path.relpath(file_path, folder_to_zip)
                zipf.write(file_path, arcname)

def process_aim_zip(aim_zip, db_directory, db_conn, cursor):
    print(f"开始处理: {aim_zip}")
    
    # 确保backup目录存在
    backup_dir = Path(db_directory) / 'backup'
    backup_dir.mkdir(exist_ok=True)
    
    # 移动旧ZIP文件到backup目录（如果有）
    old_zip_path = Path(db_directory) / aim_zip
    if old_zip_path.exists() and zipfile.is_zipfile(old_zip_path):
        new_zip_path = backup_dir / old_zip_path.name  # 使用文件名而非整个路径
        shutil.move(str(old_zip_path), str(new_zip_path))
        print(f"已移动旧ZIP文件: {old_zip_path} 到 {new_zip_path}")
    else:
        print(f"没有找到旧ZIP文件: {old_zip_path}")

    # 压缩aim_zip对应的文件夹为新的ZIP文件
    folder_to_zip = (Path(db_directory) / aim_zip).with_suffix('')
    if folder_to_zip.exists():
        zip_file_path = folder_to_zip.with_suffix('.zip')
        zip_dir(str(folder_to_zip), str(zip_file_path))
        print(f"已创建新ZIP文件: {zip_file_path}")

        # 删除原始文件夹
        shutil.rmtree(folder_to_zip)
        print(f"已删除文件夹: {folder_to_zip}")
    else:
        print(f"错误: 文件夹 {folder_to_zip} 不存在")
        return

    # 更新所有对应此aim_zip的记录，将deal_flag设置为3
    cursor.execute("UPDATE photo_move SET deal_flag=3 WHERE aim_zip=?", (aim_zip,))
    db_conn.commit()

def main(db_directory):
    db_path = str(Path(db_directory) / 'file_hashes.db')
    conn, cursor = check_and_connect_db(db_path)

    # 查询photo_move表中所有不同的aim_zip值，且deal_flag为2
    cursor.execute("SELECT DISTINCT aim_zip FROM photo_move WHERE deal_flag=2")
    aim_zips = cursor.fetchall()

    if not aim_zips:
        print("没有需要处理的ZIP文件")
        return

    for aim_zip_row in aim_zips:
        aim_zip = aim_zip_row['aim_zip']
        process_aim_zip(aim_zip, db_directory, conn, cursor)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process folders based on photo_move table, create ZIP files, and update records.")
    parser.add_argument('--db', required=True, help="Directory containing the SQLite database file (file_hashes.db)")

    args = parser.parse_args()

    db_directory = args.db

    if not os.path.isdir(db_directory):
        print(f"错误: {db_directory} 不是一个有效的目录")
        sys.exit(1)

    main(db_directory)